import torch
import datetime

from utils import *
from arch import archs
from sparseTools.utils import SparseAdjList


def adjsProcess(adjs):
    # 遍历邻接矩阵adjs的每个键值对k
    # 这句代码的整体作用是将邻接矩阵中每条边的权重重新设置为该边起始节点度的倒数。这样做的目的是进行归一化处理，使得每个节点连接的边权重之和为1，通常用于消息传递中的归一化步骤。
    # todo 这里不知道对于非 Freebase 数据集是否有作用，待检查（对于非 Freebase 数据集，使用列表存储；对于 Freebase 数据集，使用字典存储。）
    for k in adjs.keys():
        # 清除当前邻接矩阵adjs[k]的权重值。
        adjs[k].storage._value = None
        # 设置新的权重值，所有边的权重设置为平均值。
        # adjs[k] 是一个稀疏矩阵（通常是 DGLGraph 的邻接矩阵）。
        # nnz() 函数返回稀疏矩阵中非零元素的数量（number of non-zero elements）。
        # torch.ones 函数创建一个包含全为1的张量，其长度为 adjs[k].nnz()，即包含和 adjs[k] 中非零元素数量相同的1。
        # 计算邻接矩阵 adjs[k] 每行的和，结果是一个一维张量，其中每个元素代表邻接矩阵对应行的和。
        # dim=-1 表示在最后一个维度上进行求和，这里指的是行求和。
        # adjs[k].storage 是稀疏矩阵的底层存储结构，row() 返回一个张量，表示稀疏矩阵中每个非零元素的行索引。
        # 通过使用行索引张量 adjs[k].storage.row() 访问按行求和的结果，这会为每个非零元素选择其对应的行和。
        # 结果是一个张量，其长度等于 adjs[k] 中非零元素的数量，每个元素是相应非零元素所在行的和。
        # 对于每个非零元素，计算1除以该元素所在行的和。
        # 结果是一个张量，其长度为 adjs[k].nnz()，每个元素代表邻接矩阵中每条边的权重。
        # 最终，将计算得到的权重张量分配给 adjs[k].storage._value，这将更新邻接矩阵中每条边的权重。
        adjs[k].storage._value = torch.ones(adjs[k].nnz()) / adjs[k].sum(dim=-1)[adjs[k].storage.row()]

    return adjs


def neighboorAggregation(args, g, extra_metapath, tgt_type, init2sort):
    # 记录当前时间，用于计算特征传播所花费的时间。
    prop_tic = datetime.datetime.now()

    # 计算最大特征传播路径长度 max_length。
    # 如果 extra_metapath（额外的元路径）非空，则取 args.num_hops + 1 和 extra_metapath 中最长元路径长度的最大值。
    # 否则，max_length 仅设为 args.num_hops + 1。
    if len(extra_metapath):
        max_length = max(args.num_hops + 1, max([len(ele) for ele in extra_metapath]))
    else:
        max_length = args.num_hops + 1

    # 调用 hg_propagate_feat_dgl_path 函数，对图 g 的目标节点类型 tgt_type 进行特征传播。
    # 传播范围是 args.num_hops。
    # max_length 决定了传播路径的最大长度。
    # archs[args.arch][0] 表示目标数据集的第一条元路径
    # echo=False 表示在传播过程中不打印中间输出。
    g = hg_propagate_feat_dgl_path(g, tgt_type, args.num_hops, max_length, archs[args.arch][0], echo=False)

    # 初始化一个空字典 feats，用于存储目标节点类型 tgt_type 的特征。
    feats = {}
    # keys 是 g 中目标节点类型 tgt_type 上的特征键列表。
    keys = list(g.nodes[tgt_type].data.keys())
    # 打印目标节点类型和其特征键，用于调试或检查特征名称。
    print(f'For tgt {tgt_type}, feature keys {keys}')
    # 遍历目标节点的特征键 keys，将每个特征 k 从 g 的数据中移除并存储到 feats 字典中。
    # pop(k) 用于移除特征，避免将不需要的特征留在 g 中，以节省内存。
    for k in keys:
        feats[k] = g.nodes[tgt_type].data.pop(k)

    # 判断数据集是否为 DBLP、ACM 或 IMDB。
    # 如果是，data_size 记录每个特征 v 的维度大小。
    # 然后，feats 中每个特征根据 init2sort 进行重新排序。
    # 如果数据集不是这三个，assert 0 会触发异常并终止程序。
    if args.dataset in ['DBLP', 'ACM', 'IMDB']:
        data_size = {k: v.size(-1) for k, v in feats.items()}
        feats = {k: v[init2sort] for k, v in feats.items()}

    else:
        assert 0

    # 过滤 feats 字典，仅保留特征键为 archs[args.arch][0] 中的键或 tgt_type。
    feats = {k: v for k, v in feats.items() if k in archs[args.arch][0] or k == tgt_type}
    # 打印最终保留的特征键列表，便于验证或调试。
    print(
        f'After feature propagation, the list of feature keys of the target node is (metaPath) : {list(feats.keys())}')

    # 记录当前时间 prop_toc 并计算特征传播花费的时间。
    prop_toc = datetime.datetime.now()
    print(f'Time used for feat prop {prop_toc - prop_tic}')
    # 调用 gc.collect() 进行垃圾回收，释放内存。这在大规模图数据处理中尤其重要，以防内存泄漏或冗余占用。
    gc.collect()

    return feats, data_size